from sklearn.model_selection import StratifiedKFold
from sklearn.preprocessing import StandardScaler
import pandas as pd
from utils.utils import *
from sklearn.feature_selection import mutual_info_classif, VarianceThreshold, SelectKBest, mutual_info_classif
import numpy as np
class MRMRUtil:
    def __init__(self,dp,split_random_state,pic_save_path,feature_save_path=mrmr_feature_path):
        self.total_df = None
        self.feature_save_path = opj(feature_save_path,str(split_random_state))
        self.pic_save_path = opj(pic_save_path,str(split_random_state),'feature_extract','mrmr')
        self.dp = dp


    # def extract_feature(self,df,data_type,n_features_to_select=5):
    #
    #     # 先划分训练/测试集
    #     X_train, X_test, y_train, y_test, _, _ = self.dp.split_train_and_test(df.copy())
    #
    #     vt = VarianceThreshold(threshold=1e-5)
    #     X_train_vt = vt.fit_transform(X_train)
    #
    #     skb = SelectKBest(score_func=mutual_info_classif, k=('all' if X_train_vt.shape[1] < X_train_vt.shape[0] else 200))
    #     X_train_selected = skb.fit_transform(X_train_vt, y_train)
    #
    #     scaler = StandardScaler()
    #     X_train_scaled = scaler.fit_transform(X_train_selected)
    #
    #     mi = mutual_info_classif(X_train_scaled, y_train)
    #     selected_indices = np.argsort(mi)[-n_features_to_select:]
    #     selected_features = X_train.columns[selected_indices]
    #
    #     df_filtered = df.loc[:, ['Patient_ID', target_column] + selected_features.tolist()]
    #     df_filtered = df_filtered.dropna(how='all', subset=selected_features.tolist())
    #
    #     md(self.feature_save_path)
    #     df_filtered.to_csv(opj(self.feature_save_path, data_type.replace('FeatureSummary', '') + '.csv'), index=False)
    #     if self.total_df is None:
    #         self.total_df = df_filtered
    #     else:
    #         self.total_df = self.total_df.merge(df_filtered, how='left', on=['Patient_ID', target_column])


    def extract_feature(self,df,data_type,auc_threshold=0.9):

        X_train, X_test, y_train, y_test, _, _ = self.dp.split_train_and_test(df.copy())

        X, y = X_train, y_train
        feature_names = X.columns.tolist()
        skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)

        importance_dict = {f: [] for f in feature_names}

        for train_idx, val_idx in skf.split(X, y):
            X_fold, y_fold = X.iloc[train_idx], y.iloc[train_idx]

            # 方差过滤
            vt = VarianceThreshold(threshold=1e-5)
            X_vt = vt.fit_transform(X_fold)
            vt_feature_names = X_fold.columns[vt.get_support()]

            # 互信息初筛（前200或所有）
            skb = SelectKBest(score_func=mutual_info_classif,
                              k=('all' if X_vt.shape[1] < X_vt.shape[0] else 200))
            X_skb = skb.fit_transform(X_vt, y_fold)
            skb_feature_names = vt_feature_names[skb.get_support()]

            # 标准化
            scaler = StandardScaler()
            X_scaled = scaler.fit_transform(X_skb)
            X_scaled_df = pd.DataFrame(X_scaled, columns=skb_feature_names)

            # mRMR得分：使用互信息作为重要性
            mi = mutual_info_classif(X_scaled_df, y_fold, random_state=42)
            mi = np.array(mi)
            if mi.max() == mi.min():
                normalized_mi = np.zeros_like(mi)
            else:
                normalized_mi = (mi - mi.min()) / (mi.max() - mi.min() + 1e-8)

            for i, feat in enumerate(skb_feature_names):
                importance_dict[feat].append(normalized_mi[i])

        # 平均归一化得分
        mean_importance = {f: np.mean(scores) if len(scores) > 0 else 0.0 for f, scores in importance_dict.items()}
        selected_features = [f for f, score in mean_importance.items() if score >= auc_threshold]

        # 构造输出 DataFrame（包含 Patient_ID、ALN status、选中特征）
        df_filtered = df.loc[:, ['Patient_ID', target_column] + selected_features]
        df_filtered = df_filtered.dropna(how='all', subset=selected_features)

        # 保存结果
        os.makedirs(self.feature_save_path, exist_ok=True)
        df_filtered.to_csv(opj(self.feature_save_path, data_type.replace('FeatureSummary', '') + '.csv'), index=False)

        # 输出信息（可选）
        importance_df = pd.DataFrame({
            'Feature': list(mean_importance.keys()),
            'Mean_Normalized_MI': list(mean_importance.values())
        }).sort_values(by='Mean_Normalized_MI', ascending=False)

        md(self.feature_save_path)
        df_filtered.to_csv(opj(self.feature_save_path, data_type.replace('FeatureSummary', '') + '.csv'), index=False)
        if self.total_df is None:
            self.total_df = df_filtered
        else:
            self.total_df = self.total_df.merge(df_filtered, how='left', on=['Patient_ID', target_column])
        if len(selected_features) < 2:
            return True
        return False

    def save_total_df(self):
        md(self.feature_save_path)
        self.total_df.to_csv(opj(self.feature_save_path, 'total.csv'), index=False)